%% @doc Provides a gen_server-like behavior, but with some default 
%% properties thrown in, such as making fragmentron calls in response
%% to mnesia events.
%% @end

-module (gen_fragment).
-export ([ ensure_table/2, 
           start/4,
           start/5,
           start_link/4, 
           start_link/5 ]).
%-behaviour (behaviour).
-export ([ behaviour_info/1 ]).
-behaviour (gen_server).
-export ([ init/1,
           handle_call/3,
           handle_cast/2,
           handle_info/2,
           terminate/2,
           code_change/3 ]).

-record (genfrag, { module, reuse_node, state, tabs }).

%-=====================================================================-
%-                                Public                               -
%-=====================================================================-

%% @spec behaviour_info (callbacks) -> [ signature () ]
%%   signature () = { init, 1 } | { handle_call, 3 } | { handle_cast, 2 } |
%%                  { handle_info, 2 } | { terminate, 2 } | { code_change, 3 }
%% @doc The gen_server behavior signature, re-exported.  

behaviour_info (callbacks) ->
  [ { init, 1 },
    { handle_call, 3 },
    { handle_cast, 2 },
    { handle_info, 2 },
    { terminate, 2 },
    { code_change, 3 } ];
behaviour_info (_Other) ->
  undefined.

%% @spec ensure_table (TableName::atom (), TabDef::list ()) -> ok
%% @doc Convenience wrapper around fragmentron:create_table/2 which checks
%% the return value and is idempotent.
%% @end

ensure_table (TableName, TabDef) ->
  Result = case fragmentron:create_table (TableName, TabDef) of
    { atomic, ok } -> ok;
    { aborted, { already_exists, _ } } -> ok;
    { aborted, { combine_error, _, _ } } -> ok  % wtf: why not already_exists?
  end,
  ok = mnesia:wait_for_tables ([ TableName ], infinity),
  ok = mnesia:wait_for_tables ([ frag_table_name (TableName, N) 
                                 || N <- lists:seq (1, num_frags (TableName)),
                                    N > 1 ], 
                               infinity),
  Result.

%% @spec start (Module, ReuseNode::bool (), Args, Options) -> Result
%% @doc The analog to gen_server:start/3.  Takes an extra argument
%% ReuseNode which if false triggers fragmentron:remove_node/2 upon
%% termination.
%% @end

start (Module, ReuseNode, Args, Options) ->
  gen_server:start (?MODULE, [ Module, ReuseNode | Args ], Options).

%% @spec start (ServerName, Module, ReuseNode::bool (), Args, Options) -> Result
%% @doc The analog to gen_server:start/4.  Takes an extra argument
%% ReuseNode which if false triggers fragmentron:remove_node/2 upon
%% termination.
%% @end

start (ServerName, Module, ReuseNode, Args, Options) ->
  gen_server:start (ServerName, ?MODULE, [ Module, ReuseNode | Args ], Options).

%% @spec start_link (Module, ReuseNode::bool (), Args, Options) -> Result
%% @doc The analog to gen_server:start_link/3.  Takes an extra argument
%% ReuseNode which if false triggers fragmentron:remove_node/2 upon
%% termination.
%% @end

start_link (Module, ReuseNode, Args, Options) ->
  gen_server:start_link (?MODULE, [ Module, ReuseNode | Args ], Options).

%% @spec start_link (ServerName, Module, ReuseNode::bool (), Args, Options) -> Result
%% @doc The analog to gen_server:start_link/4.  Takes an extra argument
%% ReuseNode which if false triggers fragmentron:remove_node/2 upon
%% termination.
%% @end

start_link (ServerName, Module, ReuseNode, Args, Options) ->
  gen_server:start_link (ServerName, 
                         ?MODULE,
                         [ Module, ReuseNode | Args ],
                         Options).

%-=====================================================================-
%-                         gen_server callbacks                        -
%-=====================================================================-

%% @spec init (Args) -> result ()
%%   result () = { ok, Tabs::list (), State::any () } |
%%               { ok, Tabs::list (), State::any (), Timeout::integer () } |
%%               { stop, Reason::any () } |
%%               ignore
%% @doc Initialization routine.  Like Module:init/1 for gen_server, except
%% that a list of tables is returned if startup is ok, and 
%% fragmentron:add_node/2 is called for each table; and mnesia system events
%% are subscribed.
%% @end

init ([ Module, ReuseNode | Args ]) ->
  process_flag (trap_exit, true), 

  case Module:init (Args) of
    { ok, Tabs, State } ->
      lists:foreach (fun (T) -> fragmentron:add_node (T, node ()) end, Tabs),
      { ok, _ } = mnesia:subscribe (system),
      { ok, #genfrag{ module = Module, 
                      reuse_node = ReuseNode,
                      state = State,
                      tabs = Tabs } };
    { ok, Tabs, State, Timeout } ->
      lists:foreach (fun (T) -> fragmentron:add_node (T, node ()) end, Tabs),
      { ok, _ } = mnesia:subscribe (system),
      { ok, 
        #genfrag{ module = Module, 
                  reuse_node = ReuseNode,
                  state = State,
                  tabs = Tabs },
        Timeout };
    R ->
      R
  end.

%% @spec handle_call (Request, From, State) -> Result
%% @doc Just like the gen_server version.
%% @end

handle_call (Request, From, State) -> 
  wrap ((State#genfrag.module):handle_call (Request,
                                            From,
                                            State#genfrag.state),
        State).

%% @spec handle_cast (Request, State) -> Result
%% @doc Just like the gen_server version.
%% @end

handle_cast (Request, State) ->
  wrap ((State#genfrag.module):handle_cast (Request, State#genfrag.state),
        State).

%% @spec handle_info (Msg, State) -> Result
%% @doc Just like the gen_server version, except that the mnesia_system_event
%% { mnesia_down, Node } are intercepted and trigger 
%% fragmentron:rebalance_table/1 on the list of Tabs returned by 
%% Module:init/1.
%% @end

handle_info (Msg={ mnesia_system_event, { mnesia_down, _Node } }, State) -> 
  spawn (fun () -> 
           lists:foreach (fun (T) -> fragmentron:rebalance_table (T) end,
                          State#genfrag.tabs)
         end),
  wrap ((State#genfrag.module):handle_info (Msg, State#genfrag.state), State);
handle_info (Msg={ mnesia_user, { expire_old_nodes, _Nodes } }, State) -> 
  spawn (fun () -> 
           lists:foreach (fun (T) -> fragmentron:rebalance_table (T) end,
                          State#genfrag.tabs)
         end),
  wrap ((State#genfrag.module):handle_info (Msg, State#genfrag.state), State);
handle_info (Msg, State) ->
  wrap ((State#genfrag.module):handle_info (Msg, State#genfrag.state), State).

%% @spec terminate (Result, State) -> Result
%% @doc Just like the gen_server version, except that if ReuseNode is 
%% false, fragmentron:remove_node/2 is called for each table returned
%% by Module:init/1.
%% @end

terminate (Reason, State) ->
  (State#genfrag.module):terminate (Reason, State#genfrag.state),
  case State#genfrag.reuse_node of
    true ->
      ok;
    false ->
      lists:foreach (fun (T) -> fragmentron:remove_node (T, node ()) end,
                     State#genfrag.tabs)
  end.

%% @spec code_change (OldVsn, State, Extra) -> Result
%% @doc Similar to the gen_server version, but accepts either
%% { ok, State } or { ok, NewTabs, State }, where the latter 
%% replaces the set of tables monitored by the gen_fragment
%% instance.
%% @end

code_change (OldVsn, State, Extra) -> 
  case (State#genfrag.module):code_change (OldVsn,
                                           State#genfrag.state,
                                           Extra) of
    { ok, NewState } ->
      { ok, State#genfrag{ state = NewState } };
    { ok, NewTabs, NewState } ->
      { ok, State#genfrag{ tabs = NewTabs, state = NewState } }
  end.

%-=====================================================================-
%-                               Private                               -
%-=====================================================================-

frag_table_name (TableName, 1) -> TableName;
frag_table_name (TableName, FragNum) when FragNum > 1 ->
  list_to_atom (atom_to_list (TableName) ++ 
                "_frag" ++ 
                integer_to_list (FragNum)).

num_frags (TableName) ->
  FragProps = mnesia:table_info (TableName, frag_properties),
  case lists:keysearch (n_fragments, 1, FragProps) of
    { value, { n_fragments, NFrags } } -> NFrags;
    _ -> 1
  end.

wrap ({ reply, Reply, NewState }, State) ->
  { reply, Reply, State#genfrag{ state = NewState } };
wrap ({ reply, Reply, NewState, Timeout }, State) ->
  { reply, Reply, State#genfrag{ state = NewState }, Timeout };
wrap ({ noreply, NewState }, State) ->
  { noreply, State#genfrag{ state = NewState } };
wrap ({ noreply, NewState, Timeout }, State) ->
  { noreply, State#genfrag{ state = NewState }, Timeout };
wrap ({ stop, Reason, Reply, NewState }, State) ->
  { stop, Reason, Reply, State#genfrag{ state = NewState } };
wrap ({ stop, Reason, NewState }, State) ->
  { stop, Reason, State#genfrag{ state = NewState } }.
